﻿using System;
using System.Collections.Generic;
using System.Collections.ObjectModel;
using System.Linq;
using System.Text;
using System.IO;
using System.Threading;
using System.Diagnostics.CodeAnalysis;
using System.Diagnostics;

namespace Laharsub.Server.Wcf
{
    [SuppressMessage("Microsoft.Design", "CA1001:TypesThatOwnDisposableFieldsShouldBeDisposable",
        Justification = "Timer and topicsLock do not cause a memory leak because InMemoryPubSubBackend is a singleton in an app domain.")]
    public class InMemoryPubSubBackend : IPubSubBackend
    {
        int maxTopicId = 1;
        int maxMessageId = 0;
        int currentSize = 0;
        TimeSpan ttl;
        int maxSize;
        object syncRoot = new object();
        SortedList<int, PublishedTopic> topics = new SortedList<int, PublishedTopic>();
        ReaderWriterLockSlim topicsLock = new ReaderWriterLockSlim();
        Timer timer;

        public InMemoryPubSubBackend()
        {
            this.ttl = TimeSpan.FromSeconds(Laharsub.GetParamAsInt("InMemoryPubSubBackend.maxMessageTtlInS"));
            this.maxSize = Laharsub.GetParamAsInt("InMemoryPubSubBackend.maxMemoryMessageStoreInMB") * 1024 * 1024;
            this.timer = new Timer(new TimerCallback(this.PruneStaleTopics), null, (int)this.ttl.TotalMilliseconds, Timeout.Infinite);
        }

        public IAsyncResult BeginCreateTopic(AsyncCallback callback, object state)
        {
            return new CompletedAsyncResult<int>(Interlocked.Increment(ref this.maxTopicId), callback, state);
        }

        public int EndCreateTopic(IAsyncResult result)
        {
            return ((CompletedAsyncResult<int>)result).Result;
        }

        public IAsyncResult BeginPublishMessage(int topicId, string contentType, Stream body, AsyncCallback callback, object state)
        {
            if (this.currentSize > this.maxSize)
            {
                throw new InvalidOperationException("Pub/sub backend cannot accept new messages bacause the memory store quota has been reached.");
            }
            return new InMemoryPublishMessageAsyncResult(this, topicId, contentType, body, callback, state);
        }

        public int EndPublishMessage(IAsyncResult result)
        {
            return InMemoryPublishMessageAsyncResult.End(result);
        }

        public IAsyncResult BeginSubscribe(IDictionary<int, int> subscriptions, TimeSpan timeout, AsyncCallback callback, object state)
        {
            return new SubscriptionAsyncResult(this, subscriptions, timeout, callback, state);
        }

        public IEnumerable<PubSubMessage> EndSubscribe(IAsyncResult result)
        {
            return SubscriptionAsyncResult.End(result);
        }

#if CONSISTENCY
        void AssertDatabaseConsistency()
        {
            this.topicsLock.EnterReadLock();
            try
            {
                int bookmarks = 0;
                int messages = 0;
                foreach (KeyValuePair<int, PublishedTopic> entry in this.topics)
                {
                    entry.Value.Lock.EnterReadLock();
                    try
                    {
                        bookmarks += entry.Value.Bookmarks.Count;
                        messages += entry.Value.Messages.Count;
                        int last = entry.Value.Messages.Count > 0 ? entry.Value.Messages.Last().PubSubMessage.MessageId : -1;
                        if (entry.Value.Bookmarks.Any(b => b.From <= last))
                        {
                            Console.WriteLine("Consistency check failed: topicId={0}, bookmarks={1}, messages={2}",
                                entry.Key, entry.Value.Bookmarks.Count, entry.Value.Messages.Count);
                        }
                    }
                    finally
                    {
                        entry.Value.Lock.ExitReadLock();
                    }
                }
                Console.WriteLine("topics = {0}, bookmarks = {1}, messages = {2}, memory = {3}",
                    this.topics.Count, bookmarks, messages, this.currentSize);
            }
            finally
            {
                this.topicsLock.ExitReadLock();
            }
        }
#endif

        int PruneExpiredMessages(PublishedTopic publishedTopic, DateTime now)
        {
            int totalSizeReduction = 0;
            while (publishedTopic.Messages.Count > 0 && publishedTopic.Messages[0].Expire < now)
            {
                totalSizeReduction += publishedTopic.Messages[0].PubSubMessage.Body.Length;
                publishedTopic.Messages.RemoveAt(0);
            }
            return totalSizeReduction;
        }

        void PruneStaleTopics(object state)
        {
            DateTime now = DateTime.Now;
            List<int> removeCandidates = new List<int>();
            this.topicsLock.EnterUpgradeableReadLock();
            try
            {
                foreach (KeyValuePair<int, PublishedTopic> entry in this.topics)
                {
                    bool pruneExpiredMessages = false;
                    entry.Value.Lock.EnterReadLock();
                    try
                    {
                        if (this.IsTopicStale(entry.Value, now))
                        {
                            removeCandidates.Add(entry.Key);
                        }
                        else
                        {
                            pruneExpiredMessages = entry.Value.Messages.Count > 0 && entry.Value.Messages[0].Expire < now;
                        }
                    }
                    finally
                    {
                        entry.Value.Lock.ExitReadLock();
                    }
                    if (pruneExpiredMessages)
                    {
                        int totalSizeReduction;
                        entry.Value.Lock.EnterWriteLock();
                        try
                        {
                            totalSizeReduction = this.PruneExpiredMessages(entry.Value, now);
                        }
                        finally
                        {
                            entry.Value.Lock.ExitWriteLock();
                        }
                        if (totalSizeReduction > 0)
                        {
                            lock (this.syncRoot)
                            {
                                this.currentSize -= totalSizeReduction;
                            }
                        }
                    }
                }
                if (removeCandidates.Count > 0)
                {
                    this.topicsLock.EnterWriteLock();
                    try
                    {
                        foreach (int topicId in removeCandidates)
                        {
                            PublishedTopic publishedTopic = this.topics[topicId];
                            if (this.IsTopicStale(publishedTopic, now))
                            {
                                int totalSizeReduction = publishedTopic.Messages.Sum(m => m.PubSubMessage.Body.Length);
                                lock (this.syncRoot)
                                {
                                    this.currentSize -= totalSizeReduction;
                                }
                                publishedTopic.Lock.Dispose();
                                this.topics.Remove(topicId);
                            }
                        }
                    }
                    finally
                    {
                        this.topicsLock.ExitWriteLock();
                    }
                }
            }
            finally
            {
                this.topicsLock.ExitUpgradeableReadLock();
            }
            this.timer.Change((int)this.ttl.TotalMilliseconds, Timeout.Infinite);
#if CONSISTENCY
            this.AssertDatabaseConsistency();
#endif
        }

        bool IsTopicStale(PublishedTopic publishedTopic, DateTime now)
        {
            if (publishedTopic.Bookmarks.Count > 0)
            {
                return false;
            }
            if (publishedTopic.Messages.Count > 0 && publishedTopic.Messages[publishedTopic.Messages.Count - 1].Expire > now)
            {
                return false;
            }
            return true;
        }

        class PublishedTopic
        {
            public ReaderWriterLockSlim Lock { get; private set; }
            public Collection<PublishedMessage> Messages { get; private set; }
            public Collection<Bookmark> Bookmarks { get; private set; }

            public PublishedTopic()
            {
                this.Lock = new ReaderWriterLockSlim();
                this.Messages = new Collection<PublishedMessage>();
                this.Bookmarks = new Collection<Bookmark>();
            }
        }

        class PublishedMessage
        {
            public DateTime Expire { get; set; }
            public PubSubMessage PubSubMessage { get; set; }
        }

        class Bookmark
        {
            public int From { get; set; }
            public SubscriptionAsyncResult Continuation { get; set; }
            public PublishedTopic PublishedTopic { get; set; }
        }

        [SuppressMessage("Microsoft.Design", "CA1001:TypesThatOwnDisposableFieldsShouldBeDisposable",
            Justification = "Timer is disposed in the End method")]
        class SubscriptionAsyncResult : AsyncResult
        {
            static TimerCallback resumeOnTimeout = new TimerCallback(ResumeOnTimeout);
            InMemoryPubSubBackend parent;
            Timer timer;
            bool timedOut;
            PubSubMessage bookmarkTriggerMessage;
            PubSubMessage[] messages;
            Bookmark[] bookmarks;
            object syncRoot = new object();

            public SubscriptionAsyncResult(InMemoryPubSubBackend parent, IDictionary<int, int> subscriptions, TimeSpan timeout, AsyncCallback callback, object state)
                : base(callback, state)
            {
                this.parent = parent;

                foreach (int topicId in (from s in subscriptions where s.Value == 0 select s.Key).ToArray())
                {
                    subscriptions[topicId] = this.parent.maxMessageId + 1;
                }

                if (this.QueryMessages(subscriptions) || this.CreateBookmarks(subscriptions, timeout))
                {
                    this.Complete(true, null);
                }
            }

            public static IEnumerable<PubSubMessage> End(IAsyncResult result)
            {
                try
                {
                    SubscriptionAsyncResult r = AsyncResult.End<SubscriptionAsyncResult>(result);
                    return r.messages;
                }
                finally
                {
                    SubscriptionAsyncResult r1 = result as SubscriptionAsyncResult;
                    if (r1 != null && r1.timer != null)
                    {
                        r1.timer.Dispose();
                        r1.timer = null;
                    }
                }                
            }

            bool CreateBookmarks(IDictionary<int, int> subscriptions, TimeSpan timeout)
            {
                this.bookmarks = new Bookmark[subscriptions.Count];
                int bookmarkIndex = 0;
                this.parent.topicsLock.EnterUpgradeableReadLock();
                try
                {
                    foreach (KeyValuePair<int, int> subscription in subscriptions)
                    {
                        PublishedTopic publishedTopic = null;
                        bool topicExisted = true;
                        if (!this.parent.topics.TryGetValue(subscription.Key, out publishedTopic))
                        {
                            this.parent.topicsLock.EnterWriteLock();
                            try
                            {
                                if (!this.parent.topics.TryGetValue(subscription.Key, out publishedTopic))
                                {
                                    topicExisted = false;
                                    publishedTopic = new PublishedTopic();
                                    this.bookmarks[bookmarkIndex] = new Bookmark { Continuation = this, From = subscription.Value, PublishedTopic = publishedTopic };
                                    publishedTopic.Bookmarks.Add(this.bookmarks[bookmarkIndex]);
                                    bookmarkIndex++;
                                    this.parent.topics.Add(subscription.Key, publishedTopic);
                                }
                            }
                            finally
                            {
                                this.parent.topicsLock.ExitWriteLock();
                            }
                        }
                        if (topicExisted)
                        {
                            publishedTopic.Lock.EnterWriteLock();
                            try
                            {
                                int last = publishedTopic.Messages.Count > 0 ? publishedTopic.Messages.Last().PubSubMessage.MessageId : -1;
                                if (subscription.Value <= last)
                                {
                                    List<PubSubMessage> result = new List<PubSubMessage>(1);
                                    this.AddMatchingMessages(subscription.Value, publishedTopic, result);
                                    this.messages = result.ToArray();
                                    break;
                                }
                                else
                                {
                                    this.bookmarks[bookmarkIndex] = new Bookmark { Continuation = this, From = subscription.Value, PublishedTopic = publishedTopic };
                                    publishedTopic.Bookmarks.Add(this.bookmarks[bookmarkIndex]);
                                    bookmarkIndex++;
                                }
                            }
                            finally
                            {
                                publishedTopic.Lock.ExitWriteLock();
                            }
                        }
                    }
                }
                finally
                {
                    this.parent.topicsLock.ExitUpgradeableReadLock();
                }
                if (this.messages == null)
                {
                    this.timer = new Timer(resumeOnTimeout, this, timeout, TimeSpan.Zero);
                }

                return this.messages != null;
            }

            void DeleteBookmarks()
            {
                for (int i = 0; this.bookmarks != null && i < this.bookmarks.Length && this.bookmarks[i] != null; i++)
                {
                    this.bookmarks[i].PublishedTopic.Lock.EnterWriteLock();
                    try
                    {
                        this.bookmarks[i].PublishedTopic.Bookmarks.Remove(this.bookmarks[i]);
                    }
                    finally
                    {
                        this.bookmarks[i].PublishedTopic.Lock.ExitWriteLock();
                    }
                }
            }

            bool QueryMessages(IDictionary<int, int> subscriptions)
            {
                List<PubSubMessage> cummulativeResult = new List<PubSubMessage>(1);
                foreach (KeyValuePair<int, int> subscription in subscriptions)
                {
                    PublishedTopic publishedTopic = null;
                    this.parent.topicsLock.EnterReadLock();
                    try
                    {
                        if (this.parent.topics.TryGetValue(subscription.Key, out publishedTopic))
                        {
                            publishedTopic.Lock.EnterReadLock();
                            try
                            {
                                this.AddMatchingMessages(subscription.Value, publishedTopic, cummulativeResult);
                            }
                            finally
                            {
                                publishedTopic.Lock.ExitReadLock();
                            }
                        }
                    }
                    finally
                    {
                        this.parent.topicsLock.ExitReadLock();
                    }
                }
                if (cummulativeResult.Count > 0)
                {
                    this.messages = cummulativeResult.ToArray();
                }

                return this.messages != null;
            }

            void AddMatchingMessages(int from, PublishedTopic publishedTopic, List<PubSubMessage> result)
            {
                int i;
                for (i = publishedTopic.Messages.Count;
                    i > 0 && publishedTopic.Messages[i - 1].PubSubMessage.MessageId >= from;
                    i--) ;
                for (; i < publishedTopic.Messages.Count; i++)
                {
                    result.Add(publishedTopic.Messages[i].PubSubMessage);
                }
            }

            static void ResumeOnTimeout(object state)
            {
                ((SubscriptionAsyncResult)state).ResumeOnTimeout();
            }

            void ResumeOnTimeout()
            {
                bool completeSelf = false;
                lock (this.syncRoot)
                {
                    if (!this.IsCompleted && !this.timedOut && this.bookmarkTriggerMessage == null)
                    {
                        this.timedOut = true;
                        completeSelf = true;
                    }
                }
                if (completeSelf)
                {
                    this.DeleteBookmarks();
                    this.Complete(false, null);
                }
            }

            public void ResumeOnBookmark(PubSubMessage message)
            {
                bool completeSelf = false;
                lock (this.syncRoot)
                {
                    if (!this.IsCompleted && !this.timedOut && this.bookmarkTriggerMessage == null)
                    {
                        // TODO: is it possible for a bookmark to be resumed by more than one message? Should we re-query the topic?
                        this.bookmarkTriggerMessage = message;
                        completeSelf = true;
                    }
                }
                if (completeSelf)
                {
                    this.messages = new PubSubMessage[] { this.bookmarkTriggerMessage };
                    this.DeleteBookmarks();
                    this.Complete(false, null);
                }
            }
        }

        class InMemoryPublishMessageAsyncResult : PublishMessageAsyncResult
        {
            InMemoryPubSubBackend parent;

            public InMemoryPublishMessageAsyncResult(InMemoryPubSubBackend parent, int topicId, string contentType, Stream body, AsyncCallback callback, object state)
                : base(topicId, contentType, body, callback, state)
            {
                if (parent.maxTopicId < topicId)
                {
                    throw new ArgumentOutOfRangeException("topicId", "Topic does not exist.");
                }

                this.parent = parent;
                this.Start();
            }

            public override void StartPublish(bool synchronous)
            {
                PublishedTopic publishedTopic = this.PublishCore();
                this.StartResumeBookmarks(publishedTopic, synchronous);
            }

            PublishedTopic PublishCore()
            {
                PublishedTopic publishedTopic = null;
                this.parent.topicsLock.EnterUpgradeableReadLock();
                try
                {
                    bool topicExisted = true;
                    if (!this.parent.topics.TryGetValue(this.Message.TopicId, out publishedTopic))
                    {
                        this.parent.topicsLock.EnterWriteLock();
                        try
                        {
                            if (!this.parent.topics.TryGetValue(this.Message.TopicId, out publishedTopic))
                            {
                                topicExisted = false;
                                publishedTopic = new PublishedTopic();
                                this.AddMessageToTopic(publishedTopic, DateTime.Now, 0);
                                this.parent.topics.Add(this.Message.TopicId, publishedTopic);
                            }
                        }
                        finally
                        {
                            this.parent.topicsLock.ExitWriteLock();
                        }
                    }
                    if (topicExisted)
                    {
                        publishedTopic.Lock.EnterWriteLock();
                        try
                        {
                            DateTime now = DateTime.Now;
                            int totalSizeReduction = this.parent.PruneExpiredMessages(publishedTopic, now);
                            this.AddMessageToTopic(publishedTopic, now, totalSizeReduction);
                        }
                        finally
                        {
                            publishedTopic.Lock.ExitWriteLock();
                        }
                    }
                }
                finally
                {
                    this.parent.topicsLock.ExitUpgradeableReadLock();
                }

                return publishedTopic;
            }

            void AddMessageToTopic(PublishedTopic publishedTopic, DateTime now, int totalSizeReduction)
            {
                lock (this.parent.syncRoot)
                {
                    this.Message.MessageId = ++this.parent.maxMessageId;
                    this.parent.currentSize += this.Message.Body.Length - totalSizeReduction;
                }
                publishedTopic.Messages.Add(new PublishedMessage
                {
                    Expire = now.Add(this.parent.ttl),
                    PubSubMessage = this.Message
                });
            }

            void StartResumeBookmarks(PublishedTopic publishedTopic, bool synchronous)
            {
                Bookmark[] bookmarksToResume = null;
                publishedTopic.Lock.EnterReadLock();
                try
                {
                    bookmarksToResume = (from bookmark in publishedTopic.Bookmarks
                                         where bookmark.From <= this.Message.MessageId
                                         select bookmark).ToArray();
                }
                finally
                {
                    publishedTopic.Lock.ExitReadLock();
                }

                Exception callbackException = null;
                try
                {
                    foreach (Bookmark bookmark in bookmarksToResume)
                    {
                        // TODO consider scheduling bookmark resumption on a worker thread
                        bookmark.Continuation.ResumeOnBookmark(this.Message);
                    }
                }
                catch (Exception e)
                {
                    callbackException = e;
                }
                this.Complete(synchronous, callbackException);
            }
        }
    }
}
